#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# ======================================================
# @File:      : pdf_utils
# @Author     : forward_huan
# @Date       : 2025/5/13 14:59
# @Desc       :
# ======================================================
import io
import os
import re
from typing import Union

from PyPDF2 import PdfReader, PdfWriter, PdfMerger
from reportlab.pdfgen import canvas
from reportlab import rl_settings

from ...common.utils.data_utils import to_int
from ...common.utils.file_utils import checkFileExist
from ...common.utils.io_utils import closeStream
from ...common.utils.path_utils import createDirPath, pathJoin


def get_valid_page(page_str: str, page_count=-1):
    """
    获取合法的拆分页

    :param page_str: 页面数据输入，格式如"1-3, 4, 6", ","代表组数，”-”表述连续数

    :param page_count:
    :return:
    """
    if not re.fullmatch(r"[\d\-, ]*", page_str):
        raise Exception("输入数据格式不合法")
    page_str = page_str.replace(" ", "")
    if not page_str:
        return []
    page_data = page_str.split(",")
    result = []
    for page in page_data:
        if page.find("-") == -1:
            result.append([to_int(page)])
            continue
        arr = page.split("-")
        if len(arr) != 2:
            raise Exception("输入数据格式不合法")
        a, b = to_int(arr[0]), to_int(arr[1])
        if a > b:
            a, b = b, a
        result.append([i for i in range(a, b + 1)])
    for item in result:
        if page_count != -1 and (item[-1] > page_count or item[0] < 1):
            raise Exception("配置的页数不合法, [{}, {}]".format(1, page_count))
    return result


def get_reader(input_file: Union[str, PdfReader]):
    """获取reader对象"""
    if isinstance(input_file, PdfReader):
        return input_file
    checkFileExist(input_file, file_info="输入")
    return PdfReader(input_file)


def get_writer(output_file: Union[str, PdfWriter]):
    """获取reader对象"""
    if isinstance(output_file, PdfWriter):
        return output_file
    createDirPath(output_file, is_dir=False)
    return PdfWriter(output_file)


def merge_pdfs(input_files, output_file):
    """
    合并PDF文件

    :param input_files: 多个输入的pdf文件
    :param output_file: 合并后的输出文件
    :return:
    """
    merger = PdfMerger()
    checkFileExist(input_files, file_info="合并")
    for file in input_files:
        merger.append(file)

    merger.write(output_file)
    closeStream(merger)


def split_pdf(input_file: Union[str, PdfReader], output_dir, page_str):
    """
    拆分PDF文件

    :param input_file: 输入pdf文件
    :param output_dir: 输出目录
    :param page_str: 页面数据输入，格式如"1-3, 4, 6", ","代表组数，”-”表述连续数
    :return:
    """
    pdf_reader = get_reader(input_file)
    page_data = get_valid_page(page_str, len(pdf_reader.pages))
    file_name, ext = os.path.splitext(os.path.basename(input_file))
    for pages in page_data:
        pdf_writer = PdfWriter()
        for i in pages:
            pdf_writer.add_page(pdf_reader.pages[i - 1])
        output_file = pathJoin(output_dir, "{}_{}{}".format(file_name, "_".join([str(item) for item in pages]), ext))
        createDirPath(output_file, is_dir=False)
        with open(output_file, 'wb') as output:
            pdf_writer.write(output)
        closeStream(pdf_writer)
    closeStream(pdf_reader)


def extract_text(input_file: Union[str, PdfReader], page_str):
    """
    提取PDF文件文本

    :param input_file: 输入pdf文件
    :param page_str: 页面数据输入，格式如"1-3, 4, 6", ","代表组数，”-”表述连续数
    :return:
    """
    pdf_reader = get_reader(input_file)
    page_data = get_valid_page(page_str, len(pdf_reader.pages))
    result = []
    for pages in page_data:
        result.append("\n".join([pdf_reader.pages[page - 1].extract_text().strip() for page in pages]))
    closeStream(pdf_reader)
    return result


def rotate_page(input_file: Union[str, PdfReader], output_file: Union[str, PdfWriter], page_str, rotation_angle):
    """
    旋转PDF

    :param input_file: 输入pdf文件
    :param output_file: 输出pdf文件
    :param page_str: 页面数据输入，格式如"1-3, 4, 6", ","代表组数，”-”表述连续数
    :param rotation_angle: 旋转角度
    :return:
    """
    pdf_reader = get_reader(input_file)
    pdf_writer = get_writer(output_file)
    createDirPath(output_file, is_dir=False)
    page_data = get_valid_page(page_str, len(pdf_reader.pages))
    for pages in page_data:
        for page_index in pages:
            page = pdf_reader.pages[page_index - 1]
            page.rotate(rotation_angle)
            pdf_writer.add_page(page)
    with open(output_file, 'wb') as output:
        pdf_writer.write(output)
    closeStream(pdf_reader, pdf_writer)


def encrypt_pdf(input_file: Union[str, PdfReader], output_file: Union[str, PdfWriter], password):
    """
    加密PDF

    :param input_file: 输入pdf文件
    :param output_file: 输出pdf文件
    :param password: 密码
    :return:
    """
    pdf_reader = get_reader(input_file)
    pdf_writer = get_writer(output_file)
    createDirPath(output_file, is_dir=False)
    for page_num in range(len(pdf_reader.pages)):
        pdf_writer.add_page(pdf_reader.pages[page_num])
    pdf_writer.encrypt(password)
    with open(output_file, 'wb') as output:
        pdf_writer.write(output)
    closeStream(pdf_reader, pdf_writer)


def decrypt_pdf(input_file: Union[str, PdfReader], output_file: Union[str, PdfWriter], password):
    """
    解密PDF

    :param input_file: 输入pdf文件
    :param output_file: 输出pdf文件
    :param password: 密码
    :return:
    """
    pdf_reader = get_reader(input_file)
    pdf_writer = get_writer(output_file)
    createDirPath(output_file, is_dir=False)
    if pdf_reader.decrypt(password):
        for page_num in range(len(pdf_reader.pages)):
            pdf_writer.add_page(pdf_reader.pages[page_num])
        with open(output_file, 'wb') as output:
            pdf_writer.write(output)
    else:
        raise Exception("密码错误")
    closeStream(pdf_reader, pdf_writer)


def add_watermark(input_file: Union[str, PdfReader], output_file: Union[str, PdfWriter], watermark_text):
    """
    添加文字水印

    :param input_file: 输入pdf文件
    :param output_file: 输出pdf文件
    :param watermark_text: 包含水印文本的PDF
    :return:
    """
    pdf_reader = get_reader(input_file)
    pdf_writer = get_writer(output_file)
    createDirPath(output_file, is_dir=False)
    watermark = get_reader(watermark_text)
    for page_num in range(len(pdf_reader.pages)):
        page = pdf_reader.pages[page_num]
        page.merge_page(watermark.pages[0])
        pdf_writer.add_page(page)
    with open(output_file, 'wb') as output:
        pdf_writer.write(output)
    closeStream(pdf_reader, pdf_writer)


def create_page(content_data):
    """
    创建页面

    :param content_data: 创建的文本信息， [[pos, content]]
    :return:
    """

    packet = io.BytesIO()
    can = canvas.Canvas(packet)
    for pos, content in content_data:
        can.drawString(pos[0], pos[1], content)
    can.save()
    packet.seek(0)
    page = PdfReader(packet).pages[0]
    closeStream(packet)
    return page


def insert_page(input_file: Union[str, PdfReader], output_file: Union[str, PdfWriter], page_number, content_data):
    """
    插入文本

    :param input_file: 输入pdf文件
    :param output_file: 输出pdf文件
    :param page_number: 插入的页数
    :param content_data: 创建的文本信息， [[pos, content]]
    :return:
    """
    pdf_reader = get_reader(input_file)
    pdf_writer = get_writer(output_file)
    createDirPath(output_file, is_dir=False)
    for page_num in range(len(pdf_reader.pages)):
        if page_num + 1 == page_number:
            pdf_writer.add_page(create_page(content_data))
        pdf_writer.add_page(pdf_reader.pages[page_num])
    with open(output_file, 'wb') as output:
        pdf_writer.write(output)
    closeStream(pdf_reader, pdf_writer)


def delete_page(input_file: Union[str, PdfReader], output_file: Union[str, PdfWriter], page_str):
    """
    删除PDF页

    :param input_file: 输入pdf文件
    :param output_file: 输出pdf文件
    :param page_str: 页面数据输入，格式如"1-3, 4, 6", ","代表组数，”-”表述连续数
    :return:
    """
    pdf_reader = get_reader(input_file)
    pdf_writer = get_writer(output_file)
    createDirPath(output_file, is_dir=False)
    page_data = get_valid_page(page_str, len(pdf_reader.pages))
    page_data = set([page for pages in page_data for page in pages])
    for page_num in range(len(pdf_reader.pages)):
        if page_num + 1 in page_data:
            continue
        pdf_writer.add_page(pdf_reader.pages[page_num])
    with open(output_file, 'wb') as output:
        pdf_writer.write(output)
    closeStream(pdf_reader, pdf_writer)
